#include "mul.h"


template<typename T>
__global__ void cuda_mul(T* A,T* B,T* C,int rows,int cols,int total_num){
    int idx = blockIdx.x*blockDim.x+ threadIdx.x;
    if(idx<total_num){
        int row_idx = idx / cols;
        int cols_idx = idx % cols;
        C[row_idx*cols+cols_idx] = A[row_idx*cols+cols_idx] * B[row_idx*cols+cols_idx];
    }
}


template<>
__global__ void cuda_mul(half* A,half* B,half* C,int rows,int cols,int total_num){
    int idx = blockIdx.x*blockDim.x+ threadIdx.x;
    if(idx<total_num){
        int row_idx = idx / cols;
        int cols_idx = idx % cols;
        C[row_idx*cols+cols_idx] = __hmul(A[row_idx*cols+cols_idx],B[row_idx*cols+cols_idx]);
    }
}

template<typename T>
void launch_cuda_mul(T* A,T* B,T* C,int rows,int cols,int blocksize,int total_nums){
     int gridesize = (total_nums+blocksize-1)/blocksize;
     cuda_mul<T><<<gridesize,blocksize>>>(A,B,C,rows,cols,total_nums);
}


template<typename T>
void launch_cpu_mul(T* A,T* B,T* C,int rows,int cols){
    for(int row_idx=0;row_idx<rows;row_idx++){
        for(int col_idx=0;col_idx<cols;++col_idx){
            C[row_idx*cols+col_idx] = A[row_idx*cols+col_idx]*B[row_idx*cols+col_idx];
        }
    }
}

template<>
void launch_cpu_mul(half* A,half* B,half* C,int rows,int cols){
    for(int row_idx=0;row_idx<rows;row_idx++){
        for(int col_idx=0;col_idx<cols;++col_idx){
            C[row_idx*cols+col_idx] = __float2half(__half2float(A[row_idx*cols+col_idx]) * __half2float(B[row_idx*cols+col_idx]));
        }
    }
}

template void launch_cuda_mul<float>(float*,float*,float*,int,int,int,int);
template void launch_cuda_mul<half>(half*,half*,half*,int,int,int,int);
template void launch_cuda_mul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int,int,int);


template void launch_cpu_mul<float>(float*,float*,float*,int,int);
template void launch_cpu_mul<half>(half*,half*,half*,int,int);
template void launch_cpu_mul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int);
